#define DOF 1

#include "test.hpp"
#include "stencil_patterns.hpp"

int main(int argc, char* argv[])
{
    int my_pid, num_procs;
    MPI_Init(&argc, &argv);
    MPI_Comm_rank(MPI_COMM_WORLD, &my_pid);
    MPI_Comm_size(MPI_COMM_WORLD, &num_procs);

    int arg_cnt = 1;
    const std::string case_name(argv[arg_cnt++]);
    idx_t num_offset = -1;
    const int * _stencil_offset = nullptr;
    if      (case_name == "DEMO3d7" ) { num_offset =  7; _stencil_offset = stencil_offset_3d7;     }
    else if (case_name == "DEMO3d15") { num_offset = 15; _stencil_offset = stencil_offset_3d15;    }
    else if (case_name == "DEMO3d19") { num_offset = 19; _stencil_offset = stencil_offset_3d19;    }
    else if (case_name == "DEMO3d27") { num_offset = 27; _stencil_offset = stencil_offset_3d27;    }
    else if (case_name == "DEMO2d5" ) { num_offset =  5; _stencil_offset = stencil_offset_2d5i1i2; }
    else MPI_Abort(MPI_COMM_WORLD, -1111);

    // 三维划分
    // 注意这里的三维大小的第0和第1维应当与原来的SysSMG程序调转
    int case_dim    [3] = { atoi(argv[arg_cnt++]), atoi(argv[arg_cnt++]), atoi(argv[arg_cnt++]) };
    
    const idx_t glb_nblks = case_dim[0] * case_dim[1] * case_dim[2];
    idx_t my_ilower = 0, my_iupper = 0;
    {
        idx_t avg_load = glb_nblks / num_procs;
        idx_t remain   = glb_nblks - num_procs * avg_load;
        idx_t my_load = 0;
        my_ilower = avg_load * my_pid;

        if (my_pid < remain) { my_ilower += my_pid; my_load = avg_load + 1; } 
        else                 { my_ilower += remain; my_load = avg_load    ; }

        my_iupper = my_ilower + my_load - 1;
    }
    const idx_t my_nblks = my_iupper + 1 - my_ilower;

    const std::string   its_name = std::string(argv[arg_cnt++]),
                        prc_name = std::string(argv[arg_cnt++]);
    std::string config_mg_file;
	if (prc_name == "GMG") config_mg_file = std::string(argv[arg_cnt++]);
    
    if (my_pid == 0) printf("Max threads: %d\n", omp_get_max_threads());


    printf("Proc %d [%d,%d)\n", my_pid, my_ilower, my_iupper + 1);

    {
        par_CSRMatrix<idx_t, ksp_t, ksp_t, 1> A(MPI_COMM_WORLD, glb_nblks, my_ilower, my_iupper + 1,
                                                                glb_nblks, my_ilower, my_iupper + 1);
        par_Vector<idx_t, ksp_t, 1> x0(MPI_COMM_WORLD,          glb_nblks, my_ilower, my_iupper + 1);
        par_Vector<idx_t, ksp_t, 1> x(x0), b(x0), y(x0);

        x0.set_val(0.0);
        x.set_val(0.0); b.set_val(1.0); y.set_val(0.0);

        {
            long long * dist_rpt = new long long [my_nblks + 1]; dist_rpt[0] = 0;

            for (idx_t loc_1D = 0; loc_1D < my_nblks; loc_1D++) {
                const idx_t glb_1D = loc_1D + my_ilower;
                idx_t glb_3D[3], ngb_3D[3];
                glb_3D[0] =  glb_1D / (case_dim[1] * case_dim[2]);
                glb_3D[1] = (glb_1D -  case_dim[1] * case_dim[2] * glb_3D[0]) / case_dim[2];
                glb_3D[2] =  glb_1D -  case_dim[1] * case_dim[2] * glb_3D[0]  - case_dim[2] * glb_3D[1];

                idx_t cnt = 0;
                for (idx_t k = 0; k < num_offset; k++) {
                    ngb_3D[0] = glb_3D[0] + _stencil_offset[k * 3 + 0];
                    ngb_3D[1] = glb_3D[1] + _stencil_offset[k * 3 + 1];
                    ngb_3D[2] = glb_3D[2] + _stencil_offset[k * 3 + 2];

                    if (0 <= ngb_3D[0] && ngb_3D[0] < case_dim[0] &&
                        0 <= ngb_3D[1] && ngb_3D[1] < case_dim[1] &&
                        0 <= ngb_3D[2] && ngb_3D[2] < case_dim[2])
                    {
                        cnt ++;
                    }
                }
                dist_rpt[loc_1D + 1] = cnt;
            }

            for (idx_t loc_1D = 0; loc_1D < my_nblks; loc_1D++) {
                dist_rpt[loc_1D + 1] += dist_rpt[loc_1D];
            }
            
            idx_t * dist_cid = new idx_t [dist_rpt[my_nblks]];
            ksp_t * dist_val = new ksp_t [dist_rpt[my_nblks]];

            for (idx_t loc_1D = 0; loc_1D < my_nblks; loc_1D++) {
                const idx_t glb_1D = loc_1D + my_ilower;
                idx_t glb_3D[3], ngb_3D[3];
                glb_3D[0] =  glb_1D / (case_dim[1] * case_dim[2]);
                glb_3D[1] = (glb_1D -  case_dim[1] * case_dim[2] * glb_3D[0]) / case_dim[2];
                glb_3D[2] =  glb_1D -  case_dim[1] * case_dim[2] * glb_3D[0]  - case_dim[2] * glb_3D[1];

                idx_t cnt = 0;
                for (idx_t k = 0; k < num_offset; k++) {
                    ngb_3D[0] = glb_3D[0] + _stencil_offset[k * 3 + 0];
                    ngb_3D[1] = glb_3D[1] + _stencil_offset[k * 3 + 1];
                    ngb_3D[2] = glb_3D[2] + _stencil_offset[k * 3 + 2];

                    if (0 <= ngb_3D[0] && ngb_3D[0] < case_dim[0] &&
                        0 <= ngb_3D[1] && ngb_3D[1] < case_dim[1] &&
                        0 <= ngb_3D[2] && ngb_3D[2] < case_dim[2])
                    {
                        const idx_t ngb_1D = (ngb_3D[0] * case_dim[1] + ngb_3D[1]) * case_dim[2] + ngb_3D[2];
                        dist_cid[dist_rpt[loc_1D] + cnt] = ngb_1D;
                        dist_val[dist_rpt[loc_1D] + cnt] = (ngb_1D == glb_1D) ? (num_offset - 1) : -1.0;
                        cnt ++;
                    }
                }
                assert(cnt + dist_rpt[loc_1D] == dist_rpt[loc_1D + 1]);
            }

            A.set_values_distributed(dist_rpt, dist_cid, dist_val);
            delete [] dist_rpt; delete [] dist_cid; delete [] dist_val;
        }
        

        const int test_cnt = 1;
        std::vector<TEST_RECORD> records;
        for (int test = 0; test < test_cnt; test++) {
            vec_copy(x0, x);

            double ck = vec_dot<idx_t, ksp_t, double>(b, b);
            if (my_pid == 0) printf("(  b,   b) = %.15e\n", ck);

            ck = vec_dot<idx_t, ksp_t, double>(x, x);
            if (my_pid == 0) printf("(  x,   x) = %.15e\n", ck);

            A.Mult(b, y, false);
            ck = vec_dot<idx_t, ksp_t, double>(y, y);
            if (my_pid == 0) printf("(A*b, A*b) = %.27e\n", ck);

            A.Mult(x, y, false);
            ck = vec_dot<idx_t, ksp_t, double>(y, y);
            if (my_pid == 0) printf("(A*x, A*x) = %.15e\n", ck);

            TEST_CONFIG config;
            config.config_mg_file = config_mg_file;
            config.max_iter = 10;
            config.rtol = 1.0e-6;
            TEST_RECORD rec;

            buildup_solver(its_name, prc_name, config);
            setup_and_solve(A, b, x, rec);

            double true_r_norm, b_norm;
            check_residual(A, x, b, y, true_r_norm, b_norm);

            if (my_pid == 0) {
                printf("\033[1;35mtrue ||r|| = %20.16e ||r||/||b||= %20.16e\033[0m\n", 
                    true_r_norm, true_r_norm / b_norm);
                printf("Proc %d Setup, Solve costs %.6f %.6f s\n",
                    my_pid, rec.setup, rec.solve);
            }

            stat_part_times(rec);
            records.push_back(rec);

            destroy_solver();
        }// test loop

        // 输出最佳时间
        if (my_pid == 0) {
            stat_multi_runs(records);
        }
    }
    MPI_Finalize();
    return 0;
}
